#include "PCapWrapper.h"
#include <pcap.h>
#include <boost/optional.hpp>
#include <boost/shared_ptr.hpp>
#include <string>
#include <vector>
#include <windows.h>
#include "utils.h"
#include "resource.h"

namespace Broady {

PCap::PCap( void* handle )
{
	m_handle = handle;
}

PCap::~PCap( void )
{
	if( m_handle )
	{
		pcap_close( reinterpret_cast<pcap_t*>( m_handle ) );
	}
}

boost::optional<unsigned int> PCap::Dispatch( int maxPackets, HandlerProc handler, void* userData )
{
	int retVal = pcap_dispatch( reinterpret_cast<pcap_t*>( m_handle ), maxPackets, reinterpret_cast<pcap_handler>( handler ), reinterpret_cast<unsigned char*>( userData ) );

	switch( retVal )
	{
	case -1: // an error has occured
		{
			OutputMessage( "PCap Error: %s\n", pcap_geterr( reinterpret_cast<pcap_t*>( m_handle ) ) );

			return boost::optional<unsigned int>( );
		}

	case -2: // call was cancelled before any packet was processed
		{
			return boost::optional<unsigned int>( 0 );
		}

	default:
		{
			return retVal;
		}
	}
}

boost::optional<PCap::DevicePtrList> PCap::ListDevices( void )
{
	pcap_if_t* alldevs;
	char errorBuffer[ PCAP_ERRBUF_SIZE ];

	if( pcap_findalldevs( &alldevs, errorBuffer ) == -1 )
	{
		OutputMessage( "PCap Error: %s\n", errorBuffer );

		return boost::optional<PCap::DevicePtrList>( );
	}

	if( !alldevs )
	{
		OutputMessage( "No interfaces found! Make sure WinPcap is installed.\n" );

		return PCap::DevicePtrList( );
	}

	PCap::DevicePtrList lst;
	pcap_if_t* dev = alldevs;

	while( dev )
	{
		PCap::AddressList addrList;
		pcap_addr_t* pa = dev->addresses;

		while( pa )
		{
			PCap::Address a;

			if( !pa->addr )
			{
				a.m_addr = boost::optional<unsigned int>( );
			}
			else
			{
				a.m_addr = reinterpret_cast<struct sockaddr_in*>( pa->addr )->sin_addr.s_addr;
			}

			if( !pa->netmask )
			{
				a.m_netmask = boost::optional<unsigned int>( );
			}
			else
			{
				a.m_netmask = reinterpret_cast<struct sockaddr_in*>( pa->netmask )->sin_addr.s_addr;
			}

			addrList.push_back( a );

			pa = pa->next;
		}

		PCap::DevicePtr ptr = PCap::DevicePtr( new PCap::Device( dev->name, dev->description, addrList ) );
		lst.push_back( ptr );
		dev = dev->next;
	}

	pcap_freealldevs( alldevs );

	return lst;
}

boost::shared_ptr<PCap> PCap::Open( DevicePtr dev )
{
	if( !dev )
	{
		return boost::shared_ptr<PCap>( );
	}

	char errorBuffer[ PCAP_ERRBUF_SIZE ];
	unsigned int netmask;
	const char packet_filter[] = "ip and udp";
	struct bpf_program fcode;

	/* Open the adapter */
	pcap_t* adhandle = pcap_open_live(	dev->m_name.c_str( ),	/* name of the device */
										65536,					/* portion of the packet to capture. */
										1,						/* promiscuous mode */
										100,					/* read timeout */
										errorBuffer );			/* error buffer */

	if( !adhandle )
	{
		OutputMessage( "Unable to open the adapter. %s is not supported by WinPcap.\n", dev->m_name.c_str( ) );

		return boost::shared_ptr<PCap>( );
	}

	/* Check the link layer. We support only Ethernet for simplicity. */
	if( pcap_datalink( adhandle ) != DLT_EN10MB )
	{
		OutputMessage( "This program works only on Ethernet networks.\n" );
		pcap_close( adhandle );

		return boost::shared_ptr<PCap>( );
	}

	if( !dev->m_addresses.empty( ) && dev->m_addresses[0].m_netmask )
	{
		/* Retrieve the mask of the first address of the interface */
		netmask = *dev->m_addresses[0].m_netmask;
	}
	else
	{
		/* If the interface is without addresses we suppose to be in a C class network */
		netmask = 0x00FFFFFF;
	}

	/* compile the filter */
	if( pcap_compile( adhandle, &fcode, packet_filter, 1, netmask ) < 0 )
	{
		OutputMessage( "Unable to compile the packet filter. Check the syntax.\n" );
		pcap_close( adhandle );

		return boost::shared_ptr<PCap>( );
	}

	/* set the filter */
	if( pcap_setfilter( adhandle, &fcode ) < 0 )
	{
		OutputMessage( "Error setting the filter.\n" );
		pcap_close( adhandle );

		return boost::shared_ptr<PCap>( );
	}

	return boost::shared_ptr<PCap>( new PCap( adhandle ) );
}

boost::optional<bool> PCap::SetNonBlock( bool nonblock )
{
	char errorBuffer[ PCAP_ERRBUF_SIZE ];

	switch( pcap_setnonblock( reinterpret_cast<pcap_t*>( m_handle ), nonblock ? 1 : 0, errorBuffer ) )
	{
	case -1:
		{
			OutputMessage( "PCap Error: %s\n", errorBuffer );
			return boost::optional<bool>( );
		}
		
	case 0:
		{
			return false;
		}

	default:
	case 1:
		{
			return true;
		}
	}
}

PCap::Address::Address( void )
	: m_addr( )
	, m_netmask( )
{
}

PCap::Address::Address( boost::optional<unsigned int> addr, boost::optional<unsigned int> netmask )
	: m_addr( addr )
	, m_netmask( netmask )
{
}

PCap::Device::Device( void )
	: m_name( )
	, m_description( )
	, m_addresses( )
{
}

PCap::Device::Device( const char* const name, const char* const description, const std::vector<Address>& addresses )
	: m_name( name )
	, m_description( description )
	, m_addresses( addresses )
{
}

class SelectAdapter
{
public:
	static PCap::DevicePtr ShowDialog( const char* const caption, PCap::DevicePtrList alldevs );

private:
	static int CALLBACK DlgProc( HWND hDlg, UINT uMsg, WPARAM wParam, LPARAM lParam );
	static PCap::DevicePtrList s_devices;
	static PCap::DevicePtr s_retval;
};

PCap::DevicePtrList SelectAdapter::s_devices = PCap::DevicePtrList( );
PCap::DevicePtr SelectAdapter::s_retval = PCap::DevicePtr( );

int CALLBACK SelectAdapter::DlgProc( HWND hDlg, UINT uMsg, WPARAM wParam, LPARAM lParam )
{
	switch( uMsg )
	{
		case WM_INITDIALOG:
		{
			const char* caption = reinterpret_cast<const char*>( lParam );

			if( !caption )
			{
				caption = "Please select an adapter from the list.";
			}

			SetDlgItemTextA( hDlg, IDC_CAPTION, caption );

			for( PCap::DevicePtrList::const_iterator i = s_devices.begin( ), e = s_devices.end( ); i != e; i++ )
			{
				SendDlgItemMessageA( hDlg, IDC_LST, LB_ADDSTRING, 0, LPARAM( ( *i )->m_description.c_str( ) ) );
			}

			return 1;
		}

		case WM_CLOSE:
		{
			EndDialog( hDlg, IDCANCEL );
			return 1;
		}

		case WM_COMMAND:
		{
			switch( LOWORD( wParam ) )
			{
				case IDC_BTNOK:
				{
					int selIndex = static_cast<int>( SendDlgItemMessageA( hDlg, IDC_LST, LB_GETCURSEL, 0, 0 ) );

					if( selIndex == -1 || static_cast<unsigned int>( selIndex ) >= s_devices.size( ) )
					{
						return 1;
					}

					s_retval = s_devices[selIndex];
					EndDialog( hDlg, IDOK );

					return 1;
				}

				case IDC_BTNCANCEL:
				{
					EndDialog( hDlg, IDCANCEL );
					return 1;
				}

				default:
				{
					break;
				}
			}

			break;
		}

		default:
		{
			break;
		}
	}

	return 0;
}

PCap::DevicePtr SelectAdapter::ShowDialog( const char* const caption, PCap::DevicePtrList devices )
{
	s_devices = devices;
	s_retval = PCap::DevicePtr( );

	INT_PTR i = DialogBoxParamA( GetModuleHandle( 0 ), reinterpret_cast<LPCSTR>( IDD_DLGADAPTERS ), 0, reinterpret_cast<DLGPROC>( SelectAdapter::DlgProc ), LPARAM( caption ) );

	switch( i )
	{
	case -1:
		{
			OutputMessage( "Unable to create dialog %s.\n", "IDD_DLGADAPTERS" );
			break;
		}
	}

	return s_retval;
}

PCap::DevicePtr PCap::ChooseDevice( const char* const caption, const PCap::DevicePtrList& devices )
{
	return SelectAdapter::ShowDialog( caption, devices );
}

} // namespace
